import os
import torch
import numpy as np
from torch.utils.data import Dataset,DataLoader
from hparams import hparams


class VC_Dataset(Dataset):
    
    def __init__(self,para):
        
        self.path_A = para.catch_A
        self.path_B = para.catch_B
        self.n_frames = para.n_frames
        
        # 加载数据
        self.train_A = np.load(os.path.join(self.path_A,'data.npy'),allow_pickle=True).tolist()
        self.train_B = np.load(os.path.join(self.path_B,'data.npy'),allow_pickle=True).tolist()
        self.n_samples = min(len(self.train_A),len(self.train_B))
        
        # 生成随机数据对
        self.gen_random_pair_index()
        
    def gen_random_pair_index(self):
        train_data_A_idx = np.arange(len(self.train_A))
        train_data_B_idx = np.arange(len(self.train_B))
        np.random.shuffle(train_data_A_idx)
        np.random.shuffle(train_data_B_idx)
        train_data_A_idx_subset = train_data_A_idx[:self.n_samples].tolist()
        train_data_B_idx_subset = train_data_B_idx[:self.n_samples].tolist()
        self.index_pairs = [(i,j) for i,j in zip(train_data_A_idx_subset,train_data_B_idx_subset)]
     
    def __len__(self):
        return len(self.index_pairs)
        
    
    def __getitem__(self,idx):
        
        # 读取 A 和 B 的数据
        data_A = self.train_A[self.index_pairs[idx][0]]
        data_B = self.train_B[self.index_pairs[idx][1]]
        
        # 从 A 和 B 中 随机截取 n_frames 帧
        start_A = np.random.randint(len(data_A) - self.n_frames + 1)
        sub_data_A = data_A[start_A:start_A + self.n_frames]
        
        start_B = np.random.randint(len(data_B) - self.n_frames + 1)
        sub_data_B = data_B[start_B :start_B + self.n_frames]
       
        return torch.from_numpy(sub_data_A.T),torch.from_numpy(sub_data_B.T)
        
if __name__ == "__main__":
    para = hparams()
    m_Dataset= VC_Dataset(para)
    
    m_DataLoader = DataLoader(m_Dataset,batch_size = 2,shuffle = True, num_workers = 2)
    
    m_Dataset.gen_random_pair_index()
    for n_epoch in range(3):
        for i_batch, sample_batch in enumerate(m_DataLoader):
            train_A = sample_batch[0]
            train_B = sample_batch[1]
            print(train_A.shape)
            print(train_B.shape)
            if i_batch>5:
                break
            
        m_Dataset.gen_random_pair_index()
        
        
        
        